#include<ntddk.h>
//
// ioctl codes
//
#define IOCTL_APP_INIT CTL_CODE(FILE_DEVICE_UNKNOWN , 0x806 , METHOD_BUFFERED , FILE_ANY_ACCESS)
#define IOCTL_APP_EXIT CTL_CODE(FILE_DEVICE_UNKNOWN , 0x807 , METHOD_BUFFERED , FILE_ANY_ACCESS)
#define IOCTL_APP_GET  CTL_CODE(FILE_DEVICE_UNKNOWN , 0x808 , METHOD_BUFFERED , FILE_ANY_ACCESS)

//
// description of a new device
//
typedef struct _NEW_DEVICE
{
	WCHAR DeviceName[1024];
	WCHAR PDOName[1024];
} NEW_DEVICE , *PNEW_DEVICE;

//
// portotypes 
//

extern POBJECT_TYPE IoDriverObjectType;
extern NTKERNELAPI NTSTATUS ObReferenceObjectByName(IN PUNICODE_STRING ObjectName, IN ULONG Attributes, IN PACCESS_STATE PassedAccessState OPTIONAL, IN ACCESS_MASK DesiredAccess OPTIONAL, IN POBJECT_TYPE ObjectType, IN KPROCESSOR_MODE AccessMode, IN OUT PVOID ParseContext OPTIONAL, OUT PVOID *Object);

typedef NTSTATUS (*OLDDISK)( IN PDRIVER_OBJECT DriverObject , IN PDEVICE_OBJECT pdo);

//
// Global Vars
//

OLDDISK OldAddDevice; 
PKEVENT add_event;
PDRIVER_OBJECT pDisk; // driver object of disk.sys 
NEW_DEVICE new_dev;

VOID Unload( IN PDRIVER_OBJECT DriverObject );
NTSTATUS CreateCloseRoutine( IN PDEVICE_OBJECT DeviceObject , IN PIRP Irp );
NTSTATUS DispatchIoControlRoutine( IN PDEVICE_OBJECT DeviceObject , IN PIRP Irp );
NTSTATUS MyAddDevice( IN PDRIVER_OBJECT DriverObject , IN PDEVICE_OBJECT pdo );

// entry point

NTSTATUS DriverEntry( IN PDRIVER_OBJECT DriverObject , IN PUNICODE_STRING RegistryPath )
/*
   Description :
   - Create's Device \\Device\\usb_spy
   - Create's Symbolic link \\DosDevices\\susb_spy
   - get's driver object of \\Driver\Disk
   Parameters :
   -Driver Object
   -Registry Path 
   Returns:
   -STATUS_XXX
*/
{
	NTSTATUS ns = STATUS_SUCCESS;
	UNICODE_STRING usDiskDriver; // \\Driver\\Disk;
	UNICODE_STRING usDeviceName; // \\Device\\usb_spy
	UNICODE_STRING usLinkName; //  \\DosDevices\\susb_spy
	PDEVICE_OBJECT pDevObject; 

	RtlInitUnicodeString( &usDiskDriver , L"\\Driver\\Disk" );
    RtlInitUnicodeString( &usDeviceName , L"\\Device\\usb_spy" );
	RtlInitUnicodeString( &usLinkName , L"\\DosDevices\\susb_spy" );
	//
	// create device
	//
	ns = IoCreateDevice( DriverObject , 0 , &usDeviceName , FILE_DEVICE_UNKNOWN , 0 , FALSE , &pDevObject );
	if( !NT_SUCCESS(ns) )
	{
		DbgPrint("[IN DriverEntry] IoCreateDevice failed with status : %08X\n" , ns );
		return ns;
	}
	DbgPrint("[IN DriverEntry] Device Created!\n");
	//
	// create symbolic link 
	//
	ns = IoCreateSymbolicLink( &usLinkName , &usDeviceName );
	if( !NT_SUCCESS(ns) )
	{
		DbgPrint("[IN DriverEntry] IoCreateSymbolicLink failed with status : %08X\n" , ns );
		return ns;
	}
	//
	// try to get driver object of \\Driver\\Disk
	//
	ns = ObReferenceObjectByName( &usDiskDriver , OBJ_CASE_INSENSITIVE , NULL , 0 , IoDriverObjectType , KernelMode , 0 , &pDisk );
	if( !NT_SUCCESS(ns) )
	{
		DbgPrint("[IN DriverEntry] failed to get driver object with status : %08X\n" , ns );
		return ns;
	}
	DbgPrint("[IN DriverEntry] AddDevice at : %08X\n" , pDisk->DriverExtension->AddDevice);
	//
	// save original function
	//
	OldAddDevice = pDisk->DriverExtension->AddDevice;
	//
	// set handlers 
	//
	DriverObject->DriverUnload = Unload;
	DriverObject->MajorFunction[IRP_MJ_CREATE] = CreateCloseRoutine;
	DriverObject->MajorFunction[IRP_MJ_CLOSE]  = CreateCloseRoutine;
	DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = DispatchIoControlRoutine;
	return ns;
}

NTSTATUS CreateCloseRoutine( IN PDEVICE_OBJECT DeviceObject , IN PIRP Irp )
{
	Irp->IoStatus.Status = STATUS_SUCCESS;
	Irp->IoStatus.Information = 0;
    IoCompleteRequest( Irp , IO_NO_INCREMENT );
	return STATUS_SUCCESS;
}

NTSTATUS DispatchIoControlRoutine( IN PDEVICE_OBJECT DeviceObject , IN PIRP Irp )
/* Description:
   Get Information from user app.
   - Create event;
   - Hook AddDevice Routine
   - Unhook AddDevice Routine
   - Send data about new device to user app;
   Parameters:
   - Device Object;
   - Irp;
   Return:
   -STATUS_XXX
*/
{
	NTSTATUS ns = STATUS_SUCCESS;
	PIO_STACK_LOCATION pStack;
	PUCHAR Buffer;
	ULONG inBuffLen;
	ULONG outBuffLen;
	HANDLE user_event;

	pStack = IoGetCurrentIrpStackLocation( Irp );
	Buffer = Irp->AssociatedIrp.SystemBuffer;
	inBuffLen = pStack->Parameters.DeviceIoControl.InputBufferLength;
	outBuffLen = pStack->Parameters.DeviceIoControl.OutputBufferLength;

	switch( pStack->Parameters.DeviceIoControl.IoControlCode )
	{
	case IOCTL_APP_INIT:
		if( inBuffLen < sizeof(HANDLE) )
		{
			Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
			IoCompleteRequest( Irp , IO_NO_INCREMENT );
			return STATUS_INVALID_PARAMETER;
		}
		DbgPrint("[DispatchIoCOntrolRoutine] IOCTL_APP_INIT recieved!\n");
		user_event = *(PHANDLE)Buffer;
		ns = ObReferenceObjectByHandle( user_event , EVENT_MODIFY_STATE , *ExEventObjectType , UserMode , &add_event , NULL );
		if( !NT_SUCCESS(ns) )
		{
			DbgPrint("[DispatchIoControlRoutine] ObReferenceObjectByHandle failed with status : %08X\n" , ns );
			Irp->IoStatus.Status = ns;
			IoCompleteRequest( Irp , IO_NO_INCREMENT );
			return ns;
		}
		//
		// hook AddDevice Rotuine of Disk.sys 
		//
		InterlockedExchange( (PLONG)&pDisk->DriverExtension->AddDevice , (LONG)MyAddDevice);
		Irp->IoStatus.Information = 0;
		break;
	case IOCTL_APP_EXIT:
		//
		// unhook AddDevice Routine
		//
		DbgPrint("[DispatchIoControlRoutine]: IOCTL_APP_EXIT  , remove hooks.\n");
		InterlockedExchange( (PLONG)&pDisk->DriverExtension->AddDevice , (LONG)OldAddDevice);
		ObDereferenceObject( add_event );
		Irp->IoStatus.Information = 0;
		break;
	case IOCTL_APP_GET:
		if( outBuffLen < sizeof(NEW_DEVICE) )
		{
			Irp->IoStatus.Status = STATUS_INVALID_BUFFER_SIZE;
			IoCompleteRequest( Irp , IO_NO_INCREMENT );
			return STATUS_INVALID_BUFFER_SIZE;
		}
		DbgPrint("[DispatchIoControlRoutine]: IOCTL_APP_GET\n ");
		RtlCopyMemory( Buffer , &new_dev , sizeof(new_dev) );
		Irp->IoStatus.Information = sizeof(new_dev);
		break;
	default:
		ns = STATUS_INVALID_DEVICE_REQUEST;
		break;
	}
	Irp->IoStatus.Status = ns;
	IoCompleteRequest( Irp , IO_NO_INCREMENT );
	return ns;
}

NTSTATUS MyAddDevice( IN PDRIVER_OBJECT DriverObject , IN PDEVICE_OBJECT pdo )
/* Description :
   New AddDevice routine of \\Driver\\Disk
   calls IoGetDeviceProperty to get information about new device and calls original AddDevice routine;
   Parameters:
   -Driver Object
   -Device Object
   Return:
*/
{
	NTSTATUS ns;
	ULONG Size;
	memset( &new_dev , 0 , sizeof(new_dev) );
	ns = IoGetDeviceProperty( pdo , DevicePropertyFriendlyName , 1024 , &new_dev.DeviceName , &Size );
	if( NT_SUCCESS(ns) )
	{
		ns = IoGetDeviceProperty( pdo , DevicePropertyPhysicalDeviceObjectName , 1024 , &new_dev.PDOName , &Size );
		if( NT_SUCCESS(ns) )
		{


	//
	// set event for user mode
	//
	KeSetEvent( add_event , 0 , FALSE );
	DbgPrint("[MyAddDevice]: Device Name : %ws \n" , new_dev.DeviceName );
	DbgPrint("[MyAddDevice]: pdo name : %ws\n" , new_dev.PDOName );
		}
	}

	return OldAddDevice( DriverObject , pdo );
}

VOID Unload( PDRIVER_OBJECT DriverObject )
/* Description:
   Unload routine of driver.
   Delete's symbolic link and Device;
   Parameters: 
   - Driver Object;
   Return:
   -nothing
*/
{
	UNICODE_STRING usLinkName; // \\DosDevices\\susb_spy 

	RtlInitUnicodeString( &usLinkName , L"\\DosDevices\\susb_spy" );
	DbgPrint("[Unload] Deleting device and symbolic link\n");
	IoDeleteSymbolicLink( &usLinkName );
	IoDeleteDevice( DriverObject->DeviceObject );
}





















	




















